import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import os
import sys
sys.path.append("../..")
from utils.general_utils import cm2inch
import gwlsa_settings as GS
from gwlsa_settings import net_params
from utils.landslide_utils import load_df_fromfile
from utils import landslide_utils

sns.set(font='Microsoft YaHei')
sns.set(font='Microsoft YaHei')
plt.rcParams['figure.dpi'] = 300
plt.rcParams['font.sans-serif'] = ['SimHei'] #Times New Roman 或 Arial 或 SimSun 或 Segoe UI, SimHei(中文请用这个）
plt.rcParams['axes.unicode_minus'] = False  # 解决负号'-'显示为方块的问题

def plot_srcc(out_pic_name):
    train_df, val_df, test_df = load_df_fromfile(net_params['data_load_dir'], net_params['max_distance'], net_params['resolution'])
    train_val_df = pd.concat([train_df, val_df], ignore_index=False)

    df_data = train_val_df[net_params['x_column_names'] + [net_params['y_column_name']]]
    rel = df_data.corr(method='spearman')
    rel = rel.round(decimals=3)
    mask = np.zeros_like(rel)
    mask[np.triu_indices_from(mask, k=1)] = True

    #################################################################
    # SCI出图的设置,
    # 当要出19cm宽的图时,将fontsize_scale_factor设置为2.0(字体需要变大)
    # 当要出9.5cm宽的图时,将fontsize_scale_factor设置为1.0(字体保持原来大小)
    fontsize_scale_factor = 1.0
    figsize_scale_facotr = fontsize_scale_factor
    #################################################################

    cmap = sns.cubehelix_palette(start=1.5, rot=3, gamma=0.8, as_cmap=True)
    # f, ax = plt.subplots(figsize=(3.15, 2.52))
    f, ax = plt.subplots(figsize=(cm2inch(9.5*figsize_scale_facotr), cm2inch(7.6*figsize_scale_facotr)))

    sns.set(font_scale=fontsize_scale_factor)
    res = sns.heatmap(rel, ax=ax, cmap='YlGn', linewidth=0.5, annot=True,
    xticklabels=True, yticklabels=True,
                      cbar_kws={"ticks": None}, annot_kws={"fontsize": 3*fontsize_scale_factor},
                      fmt='.2f', mask=mask, square=True)
    # ax.set_title('Spearman coefficients')
    # sns.heatmap(rel, ax=ax, cmap='YlGn', linewidth=0.5, annot=True, mask=(abs(rel)< 0.5), fmt='.2f')
    ax.tick_params(labelsize=4*fontsize_scale_factor)
    cbar = ax.collections[0].colorbar
    cbar.ax.tick_params(labelsize=4*fontsize_scale_factor, axis='both', which='both', length=1)
    out_dir = GS.SAVED_PLOTS
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    plt.savefig(f'{out_dir}/{out_pic_name}', dpi=600, bbox_inches='tight')
    plt.show()

if __name__=='__main__':
    lyr_name = net_params['layer_name']
    out_pic_name = f"Figure 9.{lyr_name}因子间的spearman相关性(英文)20240921.jpg"
    landslide_utils.check_data_settings_folder()
    plot_srcc(out_pic_name)